"""
音频处理服务

Author: Video2Text Team
Date: 2024-03-20
"""
import os
import tempfile
import uuid
import shutil
from typing import Optional, Callable
import whisper
from pydub import AudioSegment

class AudioService:
    """音频处理服务"""
    
    def __init__(self):
        """初始化音频服务"""
        # 设置最大文件大小 (100MB)
        self.MAX_FILE_SIZE = 100 * 1024 * 1024  # 100MB in bytes
        
        # 支持的音频格式
        self.SUPPORTED_FORMATS = ['.mp3', '.wav', '.m4a', '.m4s', '.webm']
        
        self.model = whisper.load_model("medium")  # tiny, base, small, medium, large
        self.temp_dir = tempfile.mkdtemp()
    
    async def process_audio(self, audio_data: bytes, progress_callback: Optional[Callable] = None) -> Optional[str]:
        """处理音频数据"""
        # 输入验证
        if not audio_data:
            raise ValueError("音频数据为空")
            
        if len(audio_data) > self.MAX_FILE_SIZE:
            raise ValueError("音频文件过大")
            
        try:
            if progress_callback:
                progress_callback(10, "开始处理音频...")
            
            # 创建临时文件
            with tempfile.NamedTemporaryFile(delete=False, suffix='.m4s') as temp_file:
                temp_file.write(audio_data)
                temp_path = temp_file.name
            
            try:
                if progress_callback:
                    progress_callback(30, "音频转换中...")
                
                # 格式转换
                audio = AudioSegment.from_file(temp_path)
                processed_path = os.path.join(self.temp_dir, f"processed_{uuid.uuid4()}.wav")
                audio.export(processed_path, format="wav")
                
                if progress_callback:
                    progress_callback(60, "音频识别中...")
                
                # 音频转写
                result = self.model.transcribe(
                    processed_path,
                    language="zh",
                    task="transcribe"
                )
                
                if progress_callback:
                    progress_callback(90, "处理完成")
                
                return result["text"]
                
            finally:
                # 清理临时文件
                if os.path.exists(temp_path):
                    os.remove(temp_path)
                if os.path.exists(processed_path):
                    os.remove(processed_path)
                    
        except Exception as e:
            print(f"音频处理错误: {str(e)}")
            return None
            
    def __del__(self):
        """清理临时目录"""
        try:
            if os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
        except Exception as e:
            print(f"清理临时目录失败: {e}") 